from iotdb.Session import Session
from iotdb.utils.IoTDBConstants import TSDataType, TSEncoding, Compressor
from iotdb.utils.Tablet import Tablet

from datetime import datetime
import pandas as pd
import numpy as np

np.object = object
"""
    IotDB_utils.py
    
    @description:
        IotDB工具类
    @author:
        zhenyu
    @date:
        2024/5/28
"""


class IotDB_utils:
    tags_lst = []
    ts_path_lst = []
    measurements = []
    data_type_lst = []

    def __init__(self, host, port, user, pwd):
        self.host = str(host)
        self.port = int(port)
        self.user = str(user)
        self.password = str(pwd)
        self.maxSize = 100

        self.device = ""
        self.device_ids = device + ""
        self.storage_has_create = False

        self.session = Session(self.host, self.port, self.user, self.password, self.maxSize)
        self.session.open(enable_rpc_compression=False)
        print("登录成功")

    def create_ts(self, ts_path_lst, data_type_lst, tags_lst):
        encoding_lst = [TSEncoding.GORILLA for _ in range(len(data_type_lst))]
        compressor_lst = [Compressor.SNAPPY for _ in range(len(data_type_lst))]
        self.session.create_multi_time_series(
            ts_path_lst=ts_path_lst,
            data_type_lst=data_type_lst,
            encoding_lst=encoding_lst,
            compressor_lst=compressor_lst,
            tags_lst=tags_lst
        )

    def IotDB_insert(self, excel_path: str, device: str, device_id: str):
        if not self.storage_has_create:
            self.session.set_storage_group(device)
            self.storage_has_create = True

        # 使用pandas的read_excel函数读取Excel文件
        df = pd.read_excel(excel_path, sheet_name='Sheet1')
        
        for col in df.columns:
            self.ts_path_lst.append(device_id + "." + col)
            self.measurements.append(col)
            self.data_type_lst.append(TSDataType.DOUBLE)
            self.tags_lst.append({"dim": "dim_test"})
        print(self.ts_path_lst)
        print(self.measurements)
        print(self.tags_lst)

        if not self.session.check_time_series_exists(device_id):
            encoding_lst = [TSEncoding.GORILLA for _ in range(len(self.data_type_lst))]
            compressor_lst = [Compressor.SNAPPY for _ in range(len(self.data_type_lst))]
            self.session.create_multi_time_series(
                ts_path_lst=self.ts_path_lst,
                data_type_lst=self.data_type_lst,
                encoding_lst=encoding_lst,
                compressor_lst=compressor_lst,
                tags_lst=self.tags_lst
            )
            print("时间序列写入成功")

        data = {}
        data_nums = 0
        index = 0
        for col in df.columns:
            data[self.ts_path_lst[index]] = np.array(df[col].tolist())
            data_nums = np.array(df[col].tolist()).size
            index = index + 1

        start_time = int(datetime.timestamp(datetime.now()) * 1000)
        timestamps = np.arange(start_time, start_time + data_nums)
        df_input = pd.DataFrame(data)

        tablet = Tablet(
            device_id, self.measurements, self.data_type_lst, df_input.values, timestamps
        )
        self.session.insert_tablet(tablet)
        print("数据插入成功")


if __name__ == '__main__':
    device = "root.test2"
    device_id = device + ".uav1"
    excel_file = "D://Pycharm_Projects//GroundStation-2//src//data//fly_data.xlsx"

    username_ = "root"
    password_ = "root"
    ip = "127.0.0.1"
    port_ = 6667
    maxSize = 100

    iotdb_test = IotDB_utils(ip, port_, username_, password_)
    iotdb_test.IotDB_insert(excel_file, device, device_id)
